import numpy as np


class KKT:

    def __init__(self, env):
        self.env = env
        self.action, self.bandwidth, self.frequency = None, None, None

    def run(self, collector):
        for i in range(self.env.config.cd_episode):
            print("**************************************************************************")
            print(f"Current episode: {i}")
            # 1. 生成卸载决策（a）
            self.choose_action()
            # 2. 进行资源分配（B, f）
            self.bandwidth, self.frequency = self.ra(self.action)
            # 3. 计算当前决策的总成本
            self.env.execute(action=self.action, bandwidth=self.bandwidth, frequency=self.frequency)
            cost = self.env.calc_cost()
            t_ratio, cpu_ratio = self.env.calc_utilization()
            collector.collect(cost, sum(self.action), t_ratio, cpu_ratio)
            collector.next()
            print(f"Offloading cost: {cost}, Offloading num: {sum(self.action)}, Time Ratio:{t_ratio}, CPU Ratio:{cpu_ratio}")
            # 4. 计算资源利用率
            ratio_b = 0
            ratio_f = 0
            for g in self.env.parts_ues:
                ratio_b += sum(self.bandwidth[g]) / self.env.config.uav_bandwidth
                ratio_f += sum(self.frequency[g]) / self.env.config.uav_cpu_frequency
            print(f"Avg resource 1: {ratio_b / self.env.G}, 2: {ratio_f / self.env.G}")
        return self.action, self.bandwidth, self.frequency

    def choose_action(self):
        if self.action is None:
            self.action = np.zeros(self.env.M, dtype=int)
        else:
            cost = []
            for i in range(self.env.M):
                action = self.action.copy()
                action[i] = (0 if action[i] else 1)
                bandwidth, frequency = self.ra_without_constraint(action=action)
                ue_p = self.check_timeout(action, bandwidth, frequency)
                c = self.env.calc_cost(action=action, bandwidth=bandwidth, frequency=frequency) if len(ue_p) == 0 else 9999
                cost.append(c)
            idx = np.argmin(cost)
            self.action[idx] = (0 if self.action[idx] else 1)

    def ra(self, action):
        # 1. Calculate timeout UEs
        bandwidth, frequency = self.ra_without_constraint(action)
        ue_p = self.check_timeout(action, bandwidth, frequency)
        if len(ue_p) == 0:
            return bandwidth, frequency
        ue_p_parts = self.group_part(ue_p)
        for part in ue_p_parts:
            part_ue_p = ue_p_parts[part]
            part_ue_q = [i for i in self.env.parts_ues[part] if action[i] and i not in part_ue_p]
            # 2. Calculate nu1, nu2
            k1 = self.env.meta_trans_rate / self.env.meta_ue_loc
            k2 = self.env.meta_task_cpu / self.env.meta_ue_loc
            nu2_1 = self.env.config.uav_bandwidth * self.env.config.uav_hover_time - sum(self.env.meta_trans_rate[part_ue_p])
            # nu2_2 = sum(np.sqrt(k1[part_ue_q])) / nu2_1
            nu2_5 = sum(np.sqrt(self.env.meta_trans_rate[part_ue_p] * self.env.meta_task_cpu[part_ue_p]))
            nu2_2 = self.env.config.uav_hover_time * sum(np.sqrt(k1[part_ue_q])) * nu2_5 / nu2_1
            nu2_3 = self.env.config.uav_hover_time * sum(np.sqrt(k2[part_ue_q]))
            nu2_4 = self.env.config.uav_cpu_frequency * self.env.config.uav_hover_time - sum(self.env.meta_task_cpu[part_ue_p])
            nu2_6 = (nu2_5 ** 2) / nu2_1
            nu2 = (nu2_2 + nu2_3) / (nu2_4 - nu2_6)
            nu1 = (np.sqrt(nu2) * nu2_5 + self.env.config.uav_hover_time * sum(np.sqrt(k1[part_ue_q]))) / nu2_1
            nu1 = nu1 ** 2
            # 3. Calculate lambda, bandwidth, frequency
            lambdas = np.zeros(self.env.M, dtype=int)
            for idx in part_ue_p:
                lambdas[idx] = (np.sqrt(nu1 * k1[idx]) + np.sqrt(nu2 * k2[idx])) * self.env.meta_ue_loc[idx] / self.env.config.uav_hover_time
                lambdas[idx] = lambdas[idx] ** 2 - 1
                bandwidth[idx] = np.sqrt(k1[idx] * (1 + lambdas[idx]) / nu1)
                frequency[idx] = np.sqrt(k2[idx] * (1 + lambdas[idx]) / nu2)
            for idx in part_ue_q:
                bandwidth[idx] = np.sqrt(k1[idx] / nu1)
                frequency[idx] = np.sqrt(k2[idx] / nu2)
        return bandwidth, frequency

    def group_part(self, ue_p):
        ue_p_parts = {}
        for idx in ue_p:
            part = self.env.ues[idx].part
            if part in ue_p_parts:
                ue_p_parts[part].append(idx)
            else:
                ue_p_parts[part] = [idx]
        return ue_p_parts

    def ra_without_constraint(self, action):
        bandwidth = np.zeros(self.env.M, dtype=int)
        frequency = np.zeros(self.env.M, dtype=int)
        for g in range(self.env.G):
            off_idx = [idx for idx in self.env.parts_ues[g] if action[idx]]
            meta_trans_rate = np.sqrt(self.env.meta_trans_rate)
            meta_task_cpu = np.sqrt(self.env.meta_task_cpu)
            for ue_idx in off_idx:
                bandwidth[ue_idx] = meta_trans_rate[ue_idx] * self.env.config.uav_bandwidth / sum(
                    meta_trans_rate[off_idx])
                frequency[ue_idx] = meta_task_cpu[ue_idx] * self.env.config.uav_cpu_frequency / sum(
                    meta_task_cpu[off_idx])
        return bandwidth, frequency

    def check_timeout(self, action, bandwidth, frequency):
        with np.errstate(divide='ignore'):
            t_off = self.env.meta_trans_rate / bandwidth
            t_exe = self.env.meta_task_cpu / frequency
        for m in range(self.env.M):
            if action[m] == 0:
                t_off[m] = 0
                t_exe[m] = 0
        t_serv = t_off + t_exe
        timeout_result = (t_serv > self.env.config.uav_hover_time)
        # noinspection PyTypeChecker
        return [idx for idx, val in enumerate(timeout_result) if val]


class HeuristicsKKT:

    def __init__(self, env):
        self.env = env
        self.action, self.bandwidth, self.frequency = None, None, None

    def run(self):
        for i in range(self.env.config.cd_episode):
            print("**************************************************************************")
            print(f"Current episode: {i}")
            # 1. 生成卸载决策（a）
            action = self.action()
            # 2. 进行资源分配（B, f）
            _, bandwidth, frequency = self.ra(action)

    def action(self):
        if self.action is None:
            self.action = self.init_action()
        else:
            cost = []
            for i in range(self.env.M):
                action = self.action.copy()
                action[i] = (0 if action[i] else 1)
                bandwidth, frequency = self.ra_main(action=action)
                cost.append(self.env.calc_cost(action=action, bandwidth=bandwidth, frequency=frequency))
            idx = np.argmin(cost)
            for i in range(self.env.M):
                if self.check_swap(idx):
                    self.action[idx] = (0 if self.action[idx] else 1)
                    print(f"Swap action:{idx}={self.action[idx]}, newCost:{cost[idx]}")
                    break
                cost[idx] = np.max(cost)
                idx = np.argmin(cost)
            print(f"Swap failed, no valid action")
        return self.action

    def init_action(self):
        action = np.zeros(self.env.M, dtype=int)
        part_ues = self.env.parts_ues
        for i in range(self.env.G):
            part_action = np.random.choice([0, 1], size=len(part_ues[i]), p=[0.8, 0.2])
            for j, ue_idx in enumerate(part_ues[i]):
                action[ue_idx] = part_action[j]
        return action
        # return np.ones(self.env.M, dtype=int)

    def minus_action(self, ues_idx, action):
        for part in self.env.parts_ues:
            part_ues = [idx for idx in ues_idx if idx in part]
            min_ue = None
            min_t_loc = 10000
            for ue in self.env.ues[part_ues]:
                if ue.t_local < min_t_loc:
                    min_ue = ue
                    min_t_loc = ue.t_local
            if min_ue is not None:
                action[min_ue.idx] = 0

    def ra(self, action):
        bandwidth, frequency = self.ra_main(action)
        invalid_idx = self.check_ra(action, bandwidth, frequency)
        while len(invalid_idx) > 0:
            self.minus_action(invalid_idx, action)
            bandwidth, frequency = self.ra_main(action)
            invalid_idx = self.check_ra(action, bandwidth, frequency)
        self.action = action
        self.bandwidth = bandwidth
        self.frequency = frequency
        return action, bandwidth, frequency

    def ra_main(self, action):
        bandwidth = np.zeros(self.env.M, dtype=int)
        frequency = np.zeros(self.env.M, dtype=int)
        for g in range(self.env.G):
            off_idx = [idx for idx in self.env.parts_ues[g] if action[idx]]
            meta_trans_rate = np.sqrt(self.env.meta_trans_rate)
            meta_task_cpu = np.sqrt(self.env.meta_task_cpu)
            for ue_idx in off_idx:
                bandwidth[ue_idx] = meta_trans_rate[ue_idx] * self.env.config.uav_bandwidth[0] / sum(
                    meta_trans_rate[off_idx])
                frequency[ue_idx] = meta_task_cpu[ue_idx] * self.env.config.uav_cpu_frequency[0] / sum(
                    meta_task_cpu[off_idx])
        return bandwidth, frequency

    def check_ra(self, action, bandwidth, frequency):
        t_off = self.env.meta_trans_rate / bandwidth
        t_exe = self.env.meta_task_cpu / frequency
        for m in range(self.env.M):
            if action[m] == 0:
                t_off[m] = 0
                t_exe[m] = 0
        t_serv = t_off + t_exe
        timeout_check_result = (t_serv > self.env.config.uav_hover_time)
        # noinspection PyTypeChecker
        return [idx for idx, val in enumerate(timeout_check_result) if val]

    def check_swap(self, idx):
        action = self.action.copy()
        action[idx] = (0 if action[idx] else 1)
        bandwidth, frequency = self.ra_main(action=action)
        invalid_idx = self.check_ra(action, bandwidth, frequency)
        return len(invalid_idx) == 0